package cmm;

import org.antlr.v4.runtime.tree.TerminalNode;

import java.util.List;
import java.util.Stack;
import java.util.Vector;

public class MidCodeGenerator extends CMMBaseVisitor<String> {
    private Vector<MidCode> midCodes = new Vector<>();
    private Vector<VarSymbol> SymbolTable = new Vector<>();
    private int i = 0;
    private int number = 0;
    private int level = 0;
    private Stack<Integer> whileStack=new Stack<>();   //break语句的跳转使用

    @Override
    public String visitProgram(CMMParser.ProgramContext ctx) {
        for (int i = 0; i < ctx.getChildCount(); i++)
            visit(ctx.statement(i));
        //在程序的结尾处添加一个结束的四元式
        midCodes.add(new MidCode(number++, OpType.END, null, null, null,ctx.stop.getLine()));
        return null;
    }

    @Override
    public String visitStatement(CMMParser.StatementContext ctx) {
        visit(ctx.getChild(0));
        return null;
    }


    /**
     * 声明语句的中间代码生成
     * @param ctx
     * @return
     */
    @Override
    public String visitDecl_stat(CMMParser.Decl_statContext ctx) {
        String name;
        String type;
        int childIndex = 1;
        int childCount = ctx.getChildCount();
        List<TerminalNode>idTemp=ctx.ID();
        List<CMMParser.ExprContext>exprTemp=ctx.expr();
        do {
            if ((idTemp.contains(ctx.getChild(childIndex))&&(ctx.getChild(childIndex + 1).getText().equals(",") ||
                    ctx.getChild(childIndex + 1).getText().equals(";")))) {
                name = ctx.getChild(childIndex).getText();
                type = ctx.getChild(0).getText();
                SymbolTable.add(new VarSymbol(name, type.equals("int") ? VarType.INT : VarType.DOUBLE, level));
                midCodes.add(new MidCode(number++, OpType.DECLAR, type, null, name,ctx.start.getLine()));
            } else if (idTemp.contains(ctx.getChild(childIndex))&&ctx.getChild(childIndex + 1).getText().equals("=")&&
                    exprTemp.contains(ctx.getChild(childIndex+2))) {
                name = ctx.getChild(childIndex).getText();
                type = ctx.getChild(0).getText();
                SymbolTable.add(new VarSymbol(name, type.equals("int") ? VarType.INT : VarType.DOUBLE, level));
                midCodes.add(new MidCode(number++, OpType.DECLAR, type, null, name,ctx.start.getLine()));
                String arg1;
                arg1 = visitExpr((CMMParser.ExprContext) ctx.getChild(childIndex + 2));
                midCodes.add(new MidCode(number++, OpType.ASSIGN, arg1, null, name,ctx.start.getLine()));
            }
            childIndex+=2;
        } while (childIndex < childCount - 1);
        return null;
    }

    /**
     * 数组声明语句中间代码生成
     * @param ctx
     * @return
     */
    @Override
    public String visitArr_decl_stat(CMMParser.Arr_decl_statContext ctx) {
        String name;
        String type;
        type=ctx.getChild(0).getText().equals("int")?"intArray":"doubleArray";
        if (ctx.getChild(1).getText().equals("[")){  //int [5]a
            name=ctx.getChild(4).getText();
            String size=ctx.getChild(2).getText();
            SymbolTable.add(new VarSymbol(name,type.equals("intArray")?VarType.INTARRAY:VarType.DOUBLEARRAY,size,level));
            midCodes.add(new MidCode(number++,OpType.DECLAR,type,size,name,ctx.start.getLine()));
        }else if(ctx.getChild(2).getText().equals("[")){// int a[5]
            name=ctx.getChild(1).getText();
            String size=ctx.getChild(3).getText();
            SymbolTable.add(new VarSymbol(name,type.equals("intArray")?VarType.INTARRAY:VarType.DOUBLEARRAY,size,level));
            midCodes.add(new MidCode(number++,OpType.DECLAR,type,size,name,ctx.start.getLine()));
        }



        return null;
    }

    /**
     * 赋值语句的中间代码生成
     * @param ctx
     * @return
     */
    @Override
    public String visitAssign_stat(CMMParser.Assign_statContext ctx) {
        String result;
        String arg1;
        String arrName="";
        String index="0";
        boolean isArray=false;
        if (ctx.getChild(0).getText().equals("++")) {
            if(ctx.getChild(2).getText().equals("[")){
                arrName=ctx.getChild(1).getText();
                result =getcode_TempVar(null,arrName,ctx.start.getLine(),null);
                index=visitExpr((CMMParser.ExprContext) ctx.getChild(3));
                midCodes.add(new MidCode(number++, OpType.GET, arrName, index, result,ctx.start.getLine()));
                isArray=true;
            }else{
                result = ctx.getChild(1).getText();
            }
            midCodes.add(new MidCode(number++, OpType.PLUS, result, "1", result,ctx.start.getLine()));

        } else if (ctx.getChild(0).getText().equals("--")) {
            if(ctx.getChild(2).getText().equals("[")){
                arrName=ctx.getChild(1).getText();
                result = getcode_TempVar(null,arrName,ctx.start.getLine(),null);
                index=visitExpr((CMMParser.ExprContext) ctx.getChild(3));
                midCodes.add(new MidCode(number++, OpType.GET, arrName, index, result,ctx.start.getLine()));
                isArray=true;
            }else{
                result = ctx.getChild(1).getText();
            }
            midCodes.add(new MidCode(number++, OpType.MINS, result, "1", result,ctx.start.getLine()));
        } else {
            if(ctx.getChild(0).getChildCount()>1){
                arrName=ctx.getChild(0).getChild(0).getText();
                result=getcode_TempVar(null,arrName,ctx.start.getLine(),null);
                index=visitExpr((CMMParser.ExprContext) ctx.getChild(0).getChild(2));
                midCodes.add(new MidCode(number++,OpType.GET,arrName,index,result,ctx.start.getLine()));
                isArray=true;
            }else{
                result = ctx.getChild(0).getText();
            }

            int childCount = ctx.getChildCount();
            if (ctx.getChild(childCount-2).equals(ctx.expr())) {
                arg1 = visitExpr((CMMParser.ExprContext) ctx.getChild(childCount - 2));
                String var ;

                var=ctx.getChild(1).getText();

                switch (var) {
                    case "=":
                        MidCode FEE = new MidCode(number++, OpType.ASSIGN, arg1, null, result,ctx.start.getLine());
                        midCodes.add(FEE);
                        break;
                    case "+=":
                        MidCode FEE_p = new MidCode(number++, OpType.PLUS, result, arg1, result,ctx.start.getLine());
                        midCodes.add(FEE_p);
                        break;
                    case "*=":
                        MidCode FEE_m = new MidCode(number++, OpType.MULTI, result, arg1, result,ctx.start.getLine());
                        midCodes.add(FEE_m);
                        break;
                    case "-=":
                        MidCode FEE_mi = new MidCode(number++, OpType.MINS, result, arg1, result,ctx.start.getLine());
                        midCodes.add(FEE_mi);
                        break;
                    case "/=":
                        MidCode FEE_d = new MidCode(number++, OpType.DIV, result, arg1, result,ctx.start.getLine());
                        midCodes.add(FEE_d);
                        break;
                }
            } else {
                if (ctx.getChild(childCount - 2).getText().equals("++")) {
                    midCodes.add(new MidCode(number++, OpType.PLUS, result, "1", result,ctx.start.getLine()));
                } else {
                    midCodes.add(new MidCode(number++, OpType.MINS, result, "1", result,ctx.start.getLine()));
                }
            }
        }
        if(isArray){
            midCodes.add(new MidCode(number++,OpType.SET,arrName,index,result,ctx.start.getLine()));
        }
        return null;
    }

    @Override
    public String visitIdentifier(CMMParser.IdentifierContext ctx) {
        return super.visitIdentifier(ctx);
    }

    /**
     * if else 语句中间代码生成
     * @param ctx
     * @return
     */
    @Override
    public String visitIf_stat(CMMParser.If_statContext ctx)
    {
        Vector<Integer> backNos=new Vector<>();
        int backNo;
        if(ctx.expr().getChildCount()==3){
            OpType op=getConditionOp(ctx.getChild(2).getChild(1).getText());
            String arg1=visitCal_expr((CMMParser.Cal_exprContext) ctx.getChild(2).getChild(0));
            String arg2=visitCal_expr((CMMParser.Cal_exprContext) ctx.getChild(2).getChild(2));
            backNo=number;
            midCodes.add(new MidCode(number++,op,arg1,arg2,null,ctx.start.getLine()));
        }else{
            String arg1=visit(ctx.expr().cal_expr(0));
            backNo=number;
            midCodes.add(new MidCode(number++,OpType.JUMP_EQUAL,arg1,"0",null,ctx.start.getLine()));
        }

        visit(ctx.block(0));
        backNos.add(number);
        midCodes.add(new MidCode(number++,OpType.JUMP,null,null,null,ctx.block(0).stop.getLine()));
        int jumpNo=number;
        setMIDResultByNumber(backNo,jumpNo);
        if(ctx.getChildCount()>5){
            visit(ctx.block(1));
        }
        int endNo=number;
        for(int i=0;i<backNos.size();i++){
            setMIDResultByNumber(backNos.get(i),endNo);
        }
        return null;
    }

    /**
     * 块语句的中间代码生成
     * @param ctx
     * @return
     */
    @Override
    public String visitBlock(CMMParser.BlockContext ctx) {
        level+=1;
        midCodes.add(new MidCode(number++,OpType.CREATE,null,null,null,ctx.start.getLine()));
        for(int i=0;i<ctx.getChildCount()-2;i++){
            visit(ctx.statement(i));
        }
        midCodes.add(new MidCode(number++,OpType.CLEAR,null,null,null,ctx.stop.getLine()));
        for(int i= SymbolTable.size()-1;i>-1;i--){
            if(SymbolTable.get(i).getLevel()==level){
                SymbolTable.remove(i);
            }
        }
        level-=1;
        return null;
    }

    /**
     * while语句的中间代码生成
     * @param ctx
     * @return
     */
    @Override
    public String visitWhile_stat(CMMParser.While_statContext ctx) {
        int backNo;
        if(ctx.expr().getChildCount()==3){
            OpType op=getConditionOp(ctx.expr().getChild(1).getText());
            String arg1=visit(ctx.expr().cal_expr(0));
            String arg2=visit(ctx.expr().cal_expr(1));
            backNo=number;
            midCodes.add(new MidCode(number++,op,arg1,arg2,null,ctx.start.getLine()));
        }else{
            String arg1=visit(ctx.expr().cal_expr(0));
            backNo=number;
            midCodes.add(new MidCode(number++,OpType.JUMP_EQUAL,arg1,"0",null,ctx.start.getLine()));
        }

        visit(ctx.block());
        midCodes.add(new MidCode(number++,OpType.JUMP,null,null,null,ctx.stop.getLine()));
        setMIDResultByNumber(number-1,backNo);
        int jumpNo=number;
        if(!whileStack.empty()){
            Integer No=whileStack.pop();
            setMIDResultByNumber(No,jumpNo);
        }
        setMIDResultByNumber(backNo,jumpNo);
        return null;
    }

    @Override
    public String visitBreak_stat(CMMParser.Break_statContext ctx) {
        whileStack.push(number);
        midCodes.add(new MidCode(number++,OpType.JUMP,null,null,null,ctx.start.getLine()));
        return null;
    }

    /**
     * read语句的中间代码生成
     * @param ctx
     * @return
     */
    @Override
    public String visitRead_stat(CMMParser.Read_statContext ctx) {
        //是数组 read(a[2]);
        String arg;
        if(ctx.identifier().getChildCount()>1){
            arg=getcode_TempVar(null,ctx.identifier().getChild(0).getText(),ctx.start.getLine(),null);
            midCodes.add(new MidCode(number++,OpType.INPUT,null,null,arg,ctx.start.getLine()));
            String Index=visit(ctx.identifier().expr());
            midCodes.add(new MidCode(number++, OpType.SET,ctx.identifier().getChild(0).getText(),Index,arg,ctx.start.getLine()));
        }else{//是普通的式子
            arg=ctx.identifier().ID().getText();
            midCodes.add(new MidCode(number++, OpType.INPUT,null,null,arg,ctx.start.getLine()));
        }
        return null;
    }

    /**
     * write语句的中间代码生成
     * @param ctx
     * @return
     */
    @Override
    public String visitWrite_stat(CMMParser.Write_statContext ctx)
    {
        String arg=visit(ctx.expr());
        midCodes.add(new MidCode(number++,OpType.OUTPUT,null,null,arg,ctx.start.getLine()));
        return null;
    }

    /**
     * 表达式语句的中间代码生成
     * @param ctx
     * @return
     */
    @Override
    public String visitExpr(CMMParser.ExprContext ctx) {
        String result;
        int childCount=ctx.getChildCount();
        if(childCount==1){
            result=visitCal_expr((CMMParser.Cal_exprContext) ctx.getChild(0));
        }else{
            result=null;
        }
        return result;
    }

    /**
     * 加减因子的中间代码生成
     * @param ctx
     * @return
     */
    @Override
    public String visitCal_expr(CMMParser.Cal_exprContext ctx)
    {
        String result;
        int index;
        VarType varType=null;
        result=visitMul_expr((CMMParser.Mul_exprContext) ctx.getChild(0));
        index=2;
        if(ctx.mul_expr(0).atom(0).DOUBLE()!=null){
            varType=VarType.DOUBLE;
        }
        int childCount=ctx.getChildCount();
        while (index<childCount){
            if(ctx.getChild(index-1).getText().equals("+")){
                CMMParser.Mul_exprContext context=(CMMParser.Mul_exprContext) ctx.getChild(index);
                if(context.atom(0).DOUBLE()!=null){
                    varType=VarType.DOUBLE;
                }
                String arg2=visitMul_expr(context);
                String tempResult=getcode_TempVar(result,arg2,ctx.start.getLine(),varType);
                midCodes.add(new MidCode(number++,OpType.PLUS,result,arg2,tempResult,ctx.start.getLine()));
                result=tempResult;
                index+=2;
            }else if(ctx.getChild(index-1).getText().equals("-")){
                CMMParser.Mul_exprContext context=(CMMParser.Mul_exprContext) ctx.getChild(index);
                if(context.atom(0).DOUBLE()!=null){
                    varType=VarType.DOUBLE;
                }
                String arg2=visitMul_expr(context);
                String tempResult=getcode_TempVar(result,arg2,ctx.start.getLine(),varType);
                midCodes.add(new MidCode(number++,OpType.MINS,result,arg2,tempResult,ctx.start.getLine()));
                result=tempResult;
                index+=2;
            }
        }
        return result;
    }

    /**
     * 乘除因子的中间代码生成
     * @param ctx
     * @return
     */
    @Override
    public String visitMul_expr(CMMParser.Mul_exprContext ctx) {
        String result;
        int index;
        VarType varType=null;
        if(ctx.atom(0).DOUBLE()!=null){
            varType=VarType.DOUBLE;
        }
        if (ctx.getChild(0).getText().equals("+")){
            String arg2=visitAtom((CMMParser.AtomContext) ctx.getChild(1));
            result=getcode_TempVar(null,arg2,ctx.start.getLine(),varType);
            midCodes.add(new MidCode(number++, OpType.PLUS,"0",arg2,result,ctx.start.getLine()));
            index=3;
        }else if(ctx.getChild(0).getText().equals("-")){
            String arg2=visitAtom((CMMParser.AtomContext) ctx.getChild(1));
            result=getcode_TempVar(null,arg2,ctx.start.getLine(),varType);
            midCodes.add(new MidCode(number++, OpType.MINS,"0",arg2,result,ctx.start.getLine()));
            index=3;
        }else{
            result=visitAtom((CMMParser.AtomContext) ctx.getChild(0));
            index=2;
        }

        int childCount=ctx.getChildCount();
        while (index<childCount){
            if(ctx.getChild(index-1).getText().equals("*"))
            {   CMMParser.AtomContext context= (CMMParser.AtomContext) ctx.getChild(index);
                if(context.DOUBLE()!=null){
                    varType=VarType.DOUBLE;
                }
                String arg2=visitAtom((CMMParser.AtomContext) ctx.getChild(index));
                String  tempResult=getcode_TempVar(result,arg2,ctx.start.getLine(),varType);
                midCodes.add(new MidCode(number++, OpType.MULTI,result,arg2,tempResult,ctx.start.getLine()));
                result=tempResult;
                index+=2;
            }else if(ctx.getChild(index-1).getText().equals("/")){
                CMMParser.AtomContext context= (CMMParser.AtomContext) ctx.getChild(index);
                if(context.DOUBLE()!=null){
                    varType=VarType.DOUBLE;
                }
                String arg2=visitAtom((CMMParser.AtomContext) ctx.getChild(index));
                String tempResult=getcode_TempVar(result,arg2,ctx.start.getLine(),varType);
                midCodes.add(new MidCode(number++, OpType.DIV,result,arg2,tempResult,ctx.start.getLine()));
                result=tempResult;
                index+=2;
            }
        }
        return result;
    }

    /**
     * 最终的原子的中间代码生成
     * @param ctx
     * @return
     */
    @Override
    public String visitAtom(CMMParser.AtomContext ctx)
    {
       String result;
       String index;
       int childCount=ctx.getChild(0).getChildCount();
       if(childCount==0){
           int child=ctx.getChildCount();
           if(child==3){
               return visitExpr((CMMParser.ExprContext) ctx.getChild(1));
           }else if(child==1){
               return ctx.getChild(0).getText();
           }else if(child==4){
               index=visitExpr((CMMParser.ExprContext) ctx.getChild(2));
               result=getcode_TempVar(null,null,ctx.start.getLine(),null);
               midCodes.add(new MidCode(number++, OpType.GET,ctx.getChild(0).getText(),index,result,ctx.start.getLine()));
               return result;
           }else{
               return "";
           }
       }else if(childCount==4){
           index=visitExpr((CMMParser.ExprContext) ctx.getChild(0).getChild(2));
           String arrName=ctx.getChild(0).getChild(0).getText();
           result=getcode_TempVar(null,arrName,ctx.start.getLine(),null);
           midCodes.add(new MidCode(number++, OpType.GET,ctx.getChild(0).getChild(0).getText(),index,result,ctx.start.getLine()));
           return result;
       }else if(childCount==1) {
           return ctx.getChild(0).getText();
       }else {
           return "";
       }

    }

    /**
     * 判断操作符
     * @param op
     * @return
     */
    private OpType getConditionOp(String op){
        OpType result;
        switch (op){
            case ">":
                result=OpType.JUMP_LE;
                break;
            case "<":
                result= OpType.JUMP_GE;
                break;
            case ">=":
                result= OpType.JUMP_LOW;
                break;
            case "<=":
                result= OpType.JUMP_HIGH;
                break;
            case "==":
                result= OpType.JUMP_NE;
                break;
            case "!=":
                result= OpType.JUMP_EQUAL;
                break;
            default:
                result=OpType.JUMP;
                break;
        }
        return result;
    }

    /**
     * 获取临时变量的类型
     * @param name1
     * @param name2
     * @return
     */
    private VarType getVarType(String name1,String name2){
        VarType type1=VarType.INT;
        VarType type2=VarType.INT;
        if(name1!=null){
            for(int i=SymbolTable.size()-1;i>-1;i--){
                if(SymbolTable.get(i).getName().equals(name1)){
                    type1= SymbolTable.get(i).getType();
                }
            }
        }

        if(name2!=null){
            for(int i=SymbolTable.size()-1;i>-1;i--){
                if(SymbolTable.get(i).getName().equals(name2)){
                    type2= SymbolTable.get(i).getType();
                }
            }
        }

        if((type1.equals(VarType.INT)||type1.equals(VarType.INTARRAY))&&(type2.equals(VarType.INT)||type2.equals(VarType.INTARRAY))){
            return VarType.INT;
        }
        else{
            return VarType.DOUBLE;
        }
    }
    /**
     * get temp id
     * @param name1
     * @param name2
     * @return
     */
    private String getcode_TempVar(String name1,String name2,int lineNo,VarType varType){
        String type;
        String Result="t"+String.valueOf(i++);
        if(varType==null){
            type=getVarType(name1,name2).equals(VarType.INT)?"int":"double";
            SymbolTable.add(new VarSymbol(Result,getVarType(name1,name2),level));
        }else{
            type="double";
            SymbolTable.add(new VarSymbol(Result,varType,level));
        }

        midCodes.add(new MidCode(number++,OpType.DECLAR,type,null,Result,lineNo));
        return Result;
    }

    /**
     * Backfill according to the line number
     * @param number line number
     * @param jumpNo objective number
     */
    private void setMIDResultByNumber(int number,int jumpNo){
        for(int i=0;i<midCodes.size();i++){
            if(midCodes.get(i).getNumber()==number){
                midCodes.get(i).setResult(String.valueOf(jumpNo));
            }
        }
    }
    /**
     * get the vector of midcodes
     * @return
     */
    public Vector<MidCode> getMidCodes() {
        return midCodes;
    }
}
